package image

import (
	"bytes"
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/manifest/schema2"
	"github.com/golang/glog"
	"github.com/mholt/archiver/v3"
	"github.com/opencontainers/go-digest"
)

//Manifest manifest.json
type Manifest struct {
	Config   string   `json:"Config"`
	RepoTags []string `json:"RepoTags"`
	Layers   []string `json:"Layers"`
}

type Token struct {
	Token       string    `json:"token"`
	AccessToken string    `json:"access_token"`
	ExpiresIn   int       `json:"expires_in"`
	IssuedAt    time.Time `json:"issued_at"`
}

var (
	tokenKey = "X-Harbor-CSRF-Token"
)

type ImagePush struct {
	archivePath      string
	registryEndpoint string
	username         string
	password         string
	skipSSLVerify    bool
	tmpDir           string
	token            string
	httpClient       *http.Client
}

type HarborErrors struct {
	Errors []HarborError `json:"errors"`
}

type HarborError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

func NewImagePush(archivePath, archiveDir, registryEndpoint, username, password string, skipSSLVerify bool) *ImagePush {
	registryEndpoint = strings.TrimSuffix(registryEndpoint, "/")
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: skipSSLVerify},
	}
	return &ImagePush{
		archivePath:      archivePath,
		registryEndpoint: registryEndpoint,
		username:         username,
		password:         password,
		skipSSLVerify:    skipSSLVerify,
		tmpDir:           archiveDir,
		httpClient:       &http.Client{Transport: tr},
	}
}

func (imagePush *ImagePush) getToken() (int, error) {
	// curl -ikL -X GET -u wangxj:Harbor123456 http://10.5.26.86:8080/service/token?account=admin\&service=harbor-registry\&scope=registry:catalog:* -v
	url := fmt.Sprintf("%s/service/token?account=admin&service=harbor-registry&scope=registry:catalog:*", imagePush.registryEndpoint)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		glog.Errorf("create token request failed, err: %v", err)
		return http.StatusInternalServerError, err
	}
	req.Header.Set(tokenKey, imagePush.token)
	req.SetBasicAuth(imagePush.username, imagePush.password)
	glog.Infof("Get token %s", url)
	resp, err := imagePush.httpClient.Do(req)
	if err != nil {
		glog.Errorf("create token request failed, err: %v", err)
		return http.StatusInternalServerError, err
	}
	if resp.StatusCode != http.StatusOK {
		body, _ := ioutil.ReadAll(resp.Body)
		harborErrs := HarborErrors{}
		err = json.Unmarshal(body, &harborErrs)
		if err != nil || len(harborErrs.Errors) == 0 {
			return resp.StatusCode, fmt.Errorf("create token request failed")
		}
		return resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		glog.Errorf("read resp body failed, err: %v", err)
		return http.StatusInternalServerError, err
	}
	var token Token
	err = json.Unmarshal(body, &token)
	if err != nil {
		glog.Errorf("read resp body failed, err: %v", err)
		return http.StatusInternalServerError, err
	}
	imagePush.token = token.Token
	return resp.StatusCode, nil
}

func (imagePush *ImagePush) Push(newTag string) (int, error) {
	var err error
	var httpcode int
	httpcode, err = imagePush.getToken()
	if err != nil {
		glog.Errorf("get token failed, username: %v, newTag: %v, err: %v", imagePush.username, newTag, err)
		return httpcode, err
	}

	if !Exists(imagePush.archivePath) {
		glog.Errorf("%s not exists, username: %v, newTag: %v", imagePush.archivePath, imagePush.username, newTag)
		return http.StatusInternalServerError, fmt.Errorf("%s not exists", imagePush.archivePath)
	}

	glog.Infof("extract archive file %s to %s, username: %v, tag: %v", imagePush.archivePath, imagePush.tmpDir, imagePush.username, newTag)

	defer func() error {
		err = os.RemoveAll(imagePush.tmpDir)
		if err != nil {
			glog.Errorf("remove tmp dir %s error, %v", imagePush.tmpDir, err)
			return fmt.Errorf("remove tmp dir %s error, %v", imagePush.tmpDir, err)
		}
		return nil
	}()

	err = archiver.Unarchive(imagePush.archivePath, imagePush.tmpDir)
	if err != nil {
		glog.Errorf("unarchive failed, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
		return http.StatusInternalServerError, fmt.Errorf("unarchive failed, %+v", err)
	}
	data, err := ioutil.ReadFile(imagePush.tmpDir + "/manifest.json")
	if err != nil {
		glog.Errorf("read manifest.json failed, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
		return http.StatusInternalServerError, fmt.Errorf("read manifest.json failed, %+v", err)
	}

	var manifestObjs []*Manifest
	err = json.Unmarshal(data, &manifestObjs)
	if err != nil {
		glog.Errorf("unmarshal manifest.json failed, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
		return http.StatusInternalServerError, fmt.Errorf("unmarshal manifest.json failed, %+v", err)
	}
	image, tag, err := ParseImageAndTag(newTag)
	if err != nil {
		glog.Errorf("parse image and tag failed, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
		return http.StatusInternalServerError, err
	}
	glog.Infof("image=%s, tag=%s", image, tag)
	glog.Infof("start push image archive %s, username: %v, tag: %v, err: %+v", imagePush.archivePath, imagePush.username, newTag, err)
	for _, manifestObj := range manifestObjs {
		var layerPaths []string
		for _, layer := range manifestObj.Layers {
			layerPath := imagePush.tmpDir + "/" + layer
			httpcode, err = imagePush.pushLayer(layer, image)
			if err != nil {
				glog.Errorf("pushLayer %s Failed, username: %v, tag: %v, err: %+v", layer, imagePush.username, newTag, err)
				return httpcode, err
			}
			layerPaths = append(layerPaths, layerPath)
		}

		httpcode, err = imagePush.pushConfig(manifestObj.Config, image)
		if err != nil {
			glog.Errorf("push image config failed, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
			return httpcode, err
		}

		glog.Infof("start push manifest")
		httpcode, err = imagePush.pushManifest(layerPaths, manifestObj.Config, image, tag)
		if err != nil {
			glog.Errorf("push manifest error, username: %v, tag: %v, err: %+v", imagePush.username, newTag, err)
			return httpcode, err
		}
		glog.Infof("push manifest done, username: %v, tag: %v", imagePush.username, newTag)

	}
	glog.Infof("push image archive %s done, username: %v, tag: %v", imagePush.archivePath, imagePush.username, newTag)
	return http.StatusOK, nil
}

func (imagePush *ImagePush) checkLayerExist(file, image string) (bool, int, error) {
	hash, err := Sha256Hash(file)
	if err != nil {
		return false, http.StatusInternalServerError, err
	}
	url := fmt.Sprintf("%s/v2/%s/blobs/%s", imagePush.registryEndpoint, image,
		fmt.Sprintf("sha256:%s", hash))
	req, err := http.NewRequest("HEAD", url, nil)
	if err != nil {
		return false, http.StatusInternalServerError, err
	}
	req.Header.Set(tokenKey, imagePush.token)
	req.SetBasicAuth(imagePush.username, imagePush.password)
	glog.Infof("PUT %s", url)
	resp, err := imagePush.httpClient.Do(req)
	if err != nil {
		return false, http.StatusInternalServerError, err
	}
	if resp.StatusCode == http.StatusNotFound {
		return false, resp.StatusCode, nil
	}
	if resp.StatusCode == http.StatusOK {
		return true, resp.StatusCode, nil
	}
	body, _ := ioutil.ReadAll(resp.Body)
	harborErrs := HarborErrors{}
	err = json.Unmarshal(body, &harborErrs)
	if err != nil || len(harborErrs.Errors) == 0 {
		return false, resp.StatusCode, fmt.Errorf("head %s failed, statusCode is %d", url, resp.StatusCode)
	}
	return false, resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
}

func (imagePush *ImagePush) pushManifest(layersPaths []string, imageConfig, image, tag string) (int, error) {
	configPath := imagePush.tmpDir + "/" + imageConfig
	obj := &schema2.Manifest{}
	obj.SchemaVersion = schema2.SchemaVersion.SchemaVersion
	obj.MediaType = schema2.MediaTypeManifest
	obj.Config.MediaType = schema2.MediaTypeImageConfig
	configSize, err := GetFileSize(configPath)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	obj.Config.Size = configSize
	hash, err := Sha256Hash(configPath)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	obj.Config.Digest = digest.Digest("sha256:" + hash)
	var layerSize int64
	for _, layersPath := range layersPaths {
		layerSize, err = GetFileSize(layersPath)
		if err != nil {
			return http.StatusInternalServerError, err
		}
		hash, err = Sha256Hash(layersPath)
		if err != nil {
			return http.StatusInternalServerError, err
		}
		item := distribution.Descriptor{
			MediaType: schema2.MediaTypeUncompressedLayer,
			Size:      layerSize,
			Digest:    digest.Digest("sha256:" + hash),
		}
		obj.Layers = append(obj.Layers, item)
	}
	data, err := json.Marshal(obj)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", imagePush.registryEndpoint, image, tag)
	req, err := http.NewRequest("PUT", url, bytes.NewBuffer(data))
	if err != nil {
		return http.StatusInternalServerError, err
	}
	req.Header.Set(tokenKey, imagePush.token)
	req.SetBasicAuth(imagePush.username, imagePush.password)
	glog.Infof("PUT %s", url)
	req.Header.Set("Content-Type", schema2.MediaTypeManifest)
	resp, err := imagePush.httpClient.Do(req)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	if resp.StatusCode == http.StatusCreated {
		return resp.StatusCode, nil
	}
	glog.Errorf("put manifest failed, code is %d", resp.StatusCode)
	body, _ := ioutil.ReadAll(resp.Body)
	harborErrs := HarborErrors{}
	err = json.Unmarshal(body, &harborErrs)
	if err != nil || len(harborErrs.Errors) == 0 {
		return resp.StatusCode, fmt.Errorf("put manifest failed, code is %d", resp.StatusCode)
	}
	return resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
}

func (imagePush *ImagePush) pushConfig(imageConfig, image string) (int, error) {
	configPath := imagePush.tmpDir + "/" + imageConfig
	// check image config exists
	exist, httpCode, err := imagePush.checkLayerExist(configPath, image)
	if err != nil {
		glog.Errorf("check layer exist failed, %+v", err)
		return httpCode, err

	}
	if exist {
		glog.Infof("%s Already exist", imageConfig)
		return http.StatusOK, nil
	}

	glog.Infof("start push image config %s", imageConfig)
	url, httpCode, err := imagePush.startPushing(image)
	if err != nil {
		glog.Errorf("startPushing Error, %+v", err)
		return httpCode, err
	}
	return imagePush.chunkUpload(configPath, url)
}

func (imagePush *ImagePush) pushLayer(layer, image string) (int, error) {
	layerPath := imagePush.tmpDir + "/" + layer

	exist, httpCode, err := imagePush.checkLayerExist(layerPath, image)
	if err != nil {
		glog.Errorf("check layer exist failed, %+v", err)
		return httpCode, err
	}
	if exist {
		glog.Infof("%s Already exist", layer)
		return httpCode, nil
	}

	url, httpCode, err := imagePush.startPushing(image)
	if err != nil {
		glog.Errorf("startPushing Error, %+v", err)
		return httpCode, err
	}
	return imagePush.chunkUpload(layerPath, url)
}

func (imagePush *ImagePush) chunkUpload(file, url string) (int, error) {
	glog.Infof("push file %s to %s", file, url)
	f, err := os.Open(file)
	if err != nil {
		return http.StatusInternalServerError, err
	}
	stat, err := f.Stat()
	if err != nil {
		return http.StatusInternalServerError, err
	}
	defer f.Close()
	contentSize := stat.Size()
	chunkSize := 2097152
	index, offset := 0, 0
	buf := make([]byte, chunkSize)
	h := sha256.New()
	for {
		n, err := f.Read(buf)
		if err == io.EOF {
			break
		}
		offset = index + n
		index = offset
		glog.Infof("Pushing %s ... %.2f%s", file, (float64(offset)/float64(contentSize))*100, "%")

		chunk := buf[0:n]

		h.Write(chunk)

		if int64(offset) == contentSize {
			sum := h.Sum(nil)
			hash := hex.EncodeToString(sum)
			req, err := http.NewRequest("PUT",
				fmt.Sprintf("%s&digest=sha256:%s", url, hash), bytes.NewBuffer(chunk))
			if err != nil {
				return http.StatusInternalServerError, err
			}
			req.Header.Set(tokenKey, imagePush.token)
			req.SetBasicAuth(imagePush.username, imagePush.password)
			//glog.Infof("PUT %s", url)
			req.Header.Set("Content-Type", "application/octet-stream")
			req.Header.Set("Content-Length", fmt.Sprintf("%d", n))
			req.Header.Set("Content-Range", fmt.Sprintf("%d-%d", index, offset))
			resp, err := imagePush.httpClient.Do(req)
			if err != nil {
				return http.StatusInternalServerError, err
			}
			if resp.StatusCode != http.StatusCreated {
				body, _ := ioutil.ReadAll(resp.Body)
				harborErrs := HarborErrors{}
				err = json.Unmarshal(body, &harborErrs)
				if err != nil || len(harborErrs.Errors) == 0 {
					return resp.StatusCode, fmt.Errorf("PUT chunk layer error,code is %d", resp.StatusCode)
				}
				return resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
			}
			break
		} else {
			req, err := http.NewRequest("PATCH", url, bytes.NewBuffer(chunk))
			if err != nil {
				return http.StatusInternalServerError, err
			}
			req.Header.Set(tokenKey, imagePush.token)
			req.SetBasicAuth(imagePush.username, imagePush.password)
			req.Header.Set("Content-Type", "application/octet-stream")
			req.Header.Set("Content-Length", fmt.Sprintf("%d", n))
			req.Header.Set("Content-Range", fmt.Sprintf("%d-%d", index, offset))
			//glog.Infof("PATCH %s", url)
			resp, err := imagePush.httpClient.Do(req)
			if err != nil {
				return http.StatusInternalServerError, err
			}
			location := resp.Header.Get("Location")
			if resp.StatusCode == http.StatusAccepted && location != "" {
				url = location
			} else {
				body, _ := ioutil.ReadAll(resp.Body)
				harborErrs := HarborErrors{}
				err = json.Unmarshal(body, &harborErrs)
				if err != nil || len(harborErrs.Errors) == 0 {
					return resp.StatusCode, fmt.Errorf("PATCH chunk file error,code is %d", resp.StatusCode)
				}
				return resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
			}
		}
	}
	return http.StatusOK, nil
}

func (imagePush *ImagePush) startPushing(image string) (string, int, error) {
	url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", imagePush.registryEndpoint, image)
	req, err := http.NewRequest(http.MethodPost, url, nil)
	if err != nil {
		return "", http.StatusInternalServerError, err
	}
	req.Header.Set(tokenKey, imagePush.token)
	req.SetBasicAuth(imagePush.username, imagePush.password)
	resp, err := imagePush.httpClient.Do(req)
	if err != nil {
		return "", http.StatusInternalServerError, err
	}
	location := resp.Header.Get("Location")
	if resp.StatusCode == http.StatusAccepted && location != "" {
		return location, resp.StatusCode, nil
	}

	body, _ := ioutil.ReadAll(resp.Body)
	harborErrs := HarborErrors{}
	err = json.Unmarshal(body, &harborErrs)
	if err != nil || len(harborErrs.Errors) == 0 {
		return "", resp.StatusCode, fmt.Errorf("post %s status is %d", url, resp.StatusCode)
	}
	return "", resp.StatusCode, fmt.Errorf("%v", harborErrs.Errors[0].Message)
}

func ParseImageAndTag(repo string) (string, string, error) {
	arr := strings.Split(repo, ":")
	if len(arr) != 2 {
		return "", "", fmt.Errorf("tag is not right")
	}
	if arr[0] == "" || arr[1] == "" {
		glog.Errorf("tag is not right")
		return "", "", fmt.Errorf("tag is not right")
	}
	return arr[0], arr[1], nil
}

func Exists(path string) bool {
	_, err := os.Stat(path)
	if err != nil {
		if os.IsExist(err) {
			return true
		}
		return false
	}
	return true
}

func GetFileSize(file string) (int64, error) {
	f, err := os.Open(file)
	if err != nil {
		return 0, err
	}
	stat, err := f.Stat()
	if err != nil {
		return 0, err
	}
	defer f.Close()
	return stat.Size(), nil
}

func Sha256Hash(file string) (string, error) {
	f, err := os.Open(file)
	if err != nil {
		return "", err
	}
	defer f.Close()
	chunkSize := 65536
	buf := make([]byte, chunkSize)
	h := sha256.New()
	for {
		n, err := f.Read(buf)
		if err == io.EOF {
			break
		}
		chunk := buf[0:n]
		h.Write(chunk)
	}
	sum := h.Sum(nil)
	hash := hex.EncodeToString(sum)
	return hash, nil
}
